from langchain_community.chat_models import ChatZhipuAI
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda, RunnablePassthrough

model = ChatZhipuAI(
    model="glm-4-plus",
    temprature=0.5
)

documents = [
    Document(
        page_content="狗是伟大的伴侣，以其忠诚和友好s而闻名。",
        metadata={"source": "mammal-pets-doc"}
    ),
    Document(
        page_content="猫是独立的动物，通常喜欢自己的空间。",
        metadata={"source": "mammal-pets-doc"}
    ),
    Document(
        page_content="金鱼是初学者的流行宠物，需要相对简单的护理。",
        metadata={"source": "fish-pets-doc"}
    ),
    Document(
        page_content="鹦鹉是聪明的鸟类，能够模仿人类的语言。",
        metadata={"source": "bird-pets-doc"}
    ),
    Document(
        page_content="兔子是社交动物，需要足够的空间跳跃。",
        metadata={"source": "mammal-pets-doc"}
    )
]

embedding = ZhipuAIEmbeddings()

vector_db = Chroma.from_documents(
    documents=documents,
    embedding=embedding
)

retriever = RunnableLambda(vector_db.similarity_search).bind(k=1)

prompts = ChatPromptTemplate.from_messages([
    ("system", "请根据上下文内容来回答以下问题"),
    ("human","问题：{question} 上下文:{context}")
])

chain = {"question":RunnablePassthrough(),"context":retriever} | prompts | model

response = chain.invoke("猫是什么")
print(response.content)
